import cv2
import time
import torch
import numpy as np
import torch.nn.functional as F
import os.path as osp
import sys
import os
import threading
import queue
import datetime # 导入 datetime 模块

# 将项目根目录添加到 Python 路径，以便导入 src 模块
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), '..'))

from src.core import YAMLConfig
from src.solver import DetSolver
from src.solver.det_solver import mscoco_category2name # 导入类别名称映射

try:
    from metavision_core.event_io import EventsIterator
    from metavision_core_ml.preprocessing.viz import viz_events # 导入 viz_events 用于事件可视化
    PROPHESEE_SDK_AVAILABLE = True
except ImportError:
    print("⚠️ 未安装 Metavision SDK 或 metavision_core_ml 模块。将无法连接 Prophesee 事件相机或使用事件可视化。")
    PROPHESEE_SDK_AVAILABLE = False

class RealTimeDetector:
    def __init__(self, config_data):
        # 模拟 argparse 传递参数给 YAMLConfig
        class Args:
            def __init__(self, config_path, weights_path, ckp_resolution, confidence, use_amp=False, tuning=None):
                self.config = config_path
                self.resume = weights_path
                self.imgpath = None  # 实时检测不使用
                self.savepath = None # 实时检测不使用
                self.ckp_resolution = ckp_resolution
                self.confidence = confidence
                self.tuning = tuning
                self.amp = use_amp

        args = Args(
            config_data["config_path"],
            config_data["weights_path"],
            config_data["ckp_resolution"],
            config_data["confidence"]
        )

        cfg = YAMLConfig(
            args.config,
            resume=args.resume,
            imgpath=args.imgpath,
            savepath=args.savepath,
            ckp_resolution=args.ckp_resolution,
            confidence=args.confidence,
            use_amp=args.amp,
            tuning=args.tuning
        )
        
        self.detector = DetSolver(cfg)
        self.detector.eval() # 设置模型为评估模式

        self.cap = None # 用于传统摄像头
        self.mv_iterator = None # 用于 Prophesee 事件相机
        self.camera_type = ""

        self.frame_queue = queue.Queue(maxsize=1) # 存储处理好的帧，只保留最新一帧
        self.running = False # 控制相机捕获线程的标志
        self.camera_thread = None

        self.fps = 0
        self.frame_count = 0
        self.start_time = time.time()
        self.ckp_resolution = config_data["ckp_resolution"]
        self.confidence = config_data["confidence"]
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.camera_source = config_data["camera_source"]
        self.frame_width = config_data["frame_width"]
        self.frame_height = config_data["frame_height"]
        self.save_video_path = config_data.get("save_video_path", None) # 新增保存视频路径
        self.video_writer = None # 新增 video_writer
        self.target_processing_fps = config_data.get("target_processing_fps", 30) # 新增目标处理帧率
        
        # 将模型移动到指定的设备
        if self.detector.model is not None:
            self.detector.model.to(self.device)
        if self.detector.ema is not None and self.detector.ema.module is not None:
            self.detector.ema.module.to(self.device)

    def setup_camera(self, src=0, width=640, height=480):
        """初始化相机并显示连接状态"""
        print("🔄 正在尝试连接相机...")
        
        if isinstance(src, int) and PROPHESEE_SDK_AVAILABLE and src == 0: # 假设 0 是默认 Prophesee 相机
            try:
                self.mv_iterator = EventsIterator(input_path="") # 移除 max_events
                self.camera_type = "prophesee"
                print(f"📷 尝试连接 Prophesee 事件相机 (源: {src})")
            except Exception as e:
                print(f"❌ 无法连接 Prophesee 事件相机: {e}")
                # 如果 Prophesee 连接失败，且源为整数 (如 0)，尝试回退到传统相机
                if isinstance(src, int):
                    print(f"🔄 Prophesee 事件相机连接失败，尝试将源 {src} 作为传统相机连接...")
                    self.cap = cv2.VideoCapture(src)
                    if self.cap.isOpened():
                        self.camera_type = "traditional"
                        print(f"✅ 成功连接传统相机 (源: {src})")
                    else:
                        self.camera_type = "" # 传统相机也连接失败
                        print(f"❌ 无法连接传统相机 (源: {src})")
                else:
                    self.camera_type = "" # 非整数源，无回退
        elif isinstance(src, str) and PROPHESEE_SDK_AVAILABLE and (src.endswith(".raw") or src.startswith("file://")):
            try:
                self.mv_iterator = EventsIterator(input_path=src, delta_t=10000) # 为文件读取设置 50ms 的 delta_t
                self.camera_type = "prophesee_file"
                print(f"📂 尝试读取 Prophesee 事件文件 (路径: {src})")
            except Exception as e:
                print(f"❌ 无法打开 Prophesee 事件文件: {e}")
                self.camera_type = ""
        elif isinstance(src, int) or (isinstance(src, str) and (src.startswith("rtsp://") or src.startswith("http://"))):
            self.cap = cv2.VideoCapture(src)
            self.camera_type = "traditional"
            # 显示连接参数
            if isinstance(src, int):
                print(f"📷 尝试访问设备索引号: {src} (0=默认摄像头)")
            else:
                print(f"🌐 尝试连接网络流: {src}")

        if self.camera_type == "traditional":
            # 设置分辨率
            self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, width)
            self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, height)
            
            # 检查连接状态
            if not self.cap.isOpened():
                error_msg = f"❌ 无法连接相机 (源: {src})"
                if isinstance(src, int):
                    error_msg += "\n可能原因:\n1. 摄像头未正确连接\n2. 驱动未安装\n3. 被其他程序占用"
                else:
                    error_msg += "\n可能原因:\n1. 网络地址错误\n2. 需要认证\n3. 端口被阻止"
                raise RuntimeError(error_msg)
            
            # 连接成功提示
            print("✅ 相机连接成功!")
            print(f"📏 分辨率设置为: {width}x{height}")
            
            # 显示实际分辨率（可能与设置不同）
            actual_width = int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH))
            actual_height = int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
            if (actual_width, actual_height) != (width, height):
                print(f"⚠️ 注意: 实际分辨率为 {actual_width}x{actual_height} (可能不支持设置的分辨率)")
            
            # 获取帧率用于视频保存，如果获取不到，使用目标处理帧率
            # fps = self.cap.get(cv2.CAP_PROP_FPS) # 移除或注释掉，直接使用目标帧率
            # if fps == 0:
            #     print("⚠️ 无法获取相机帧率，将使用默认值 30 FPS 进行视频保存。")
            #     fps = 30.0

        elif self.camera_type == "prophesee" or self.camera_type == "prophesee_file":
            print("✅ Prophesee 相机/文件连接成功!")
            # Prophesee 相机分辨率通常由其内部决定，此处不设置
            print(f"📏 事件流分辨率 (将转换为 {width}x{height} 用于检测): {self.frame_width}x{self.frame_height}")
            # fps = 30.0 # 移除或注释掉，直接使用目标帧率
        else:
            raise RuntimeError(f"❌ 无法连接相机或文件 (源: {src})。请检查源类型或 Metavision SDK 安装。")

        # 如果指定了保存路径，则初始化视频写入器
        if self.save_video_path:
            if not self.save_video_path.lower().endswith(".mp4"):
                raise ValueError(f"❌ 视频保存路径必须是 .mp4 格式: {self.save_video_path}")
            
            fourcc = cv2.VideoWriter_fourcc(*'MP4V') # 尝试使用 MP4V 编码器，通常兼容性更好
            # 将目标处理帧率除以 2，以解决视频播放速度过快的问题
            self.video_writer = cv2.VideoWriter(self.save_video_path, fourcc, self.target_processing_fps, (width, height)) # 使用 target_processing_fps
            if not self.video_writer.isOpened():
                raise RuntimeError(f"❌ 无法创建视频文件写入器: {self.save_video_path}")
            print(f"💾 视频将保存到: {self.save_video_path}")

    def _camera_capture_loop(self):
        """相机捕获和事件到帧转换的循环，在新线程中运行"""
        mv_iterator_instance = None
        if self.camera_type == "prophesee" or self.camera_type == "prophesee_file":
            try:
                # 确保获取到一个真正的迭代器对象
                mv_iterator_instance = iter(self.mv_iterator)
            except TypeError as e:
                print(f"❌ [相机线程] 无法从 EventsIterator 获取迭代器: {e}")
                self.running = False
                return

        while self.running:
            frame = None
            if self.camera_type == "traditional":
                ret, frame = self.cap.read()
                if not ret:
                    print("⚠️ [相机线程] 无法获取帧，可能相机已断开连接")
                    self.running = False # 停止线程
                    break
            elif self.camera_type == "prophesee" or self.camera_type == "prophesee_file":
                try:
                    events = next(mv_iterator_instance, None) # 使用获取到的迭代器
                    if events is None: # 如果事件流结束 (例如文件读取完毕)
                        if self.camera_type == "prophesee_file":
                            print("⏹️ [相机线程] 事件流已结束。")
                        else: # 对于实时相机，这可能意味着连接问题或相机停止流
                            print("❌ [相机线程] Prophesee 实时事件流意外结束或断开连接。")
                        self.running = False
                        break
                    frame = self._events_to_frame(events, self.frame_width, self.frame_height)
                    if frame is None:
                        print("⚠️ [相机线程] 无法将事件转换为图像帧。")
                        # 如果 frame 异常为 None，短暂等待避免忙循环
                        time.sleep(0.01)
                        continue
                except Exception as e:
                    print(f"❌ [相机线程] 处理 Prophesee 事件时发生错误: {e}")
                    self.running = False
                    break
            else:
                print("❌ [相机线程] 未知的相机类型或未成功初始化相机。")
                self.running = False # 停止线程
                break

            if frame is not None:
                # 清空队列，只保留最新一帧
                if not self.frame_queue.empty():
                    try:
                        self.frame_queue.get_nowait() # 尝试取出旧帧
                    except queue.Empty:
                        pass
                self.frame_queue.put(frame) # 放入新帧

    def run(self):
        """主循环"""
        try:
            print("\n🎥 开始实时检测 (按Q键退出)...")
            self.running = True
            self.camera_thread = threading.Thread(target=self._camera_capture_loop, daemon=True)
            self.camera_thread.start()

            # 初始化帧率计算变量，以便在循环中持续更新和打印
            self.fps = 0
            self.frame_count = 0
            self.start_time = time.time()

            while self.running:
                frame_start_time = time.time() # 记录当前帧处理开始时间

                try:
                    # 从队列中获取最新帧，超时1秒，防止无限等待
                    frame = self.frame_queue.get(timeout=1)
                except queue.Empty:
                    # 如果队列为空，可能是相机线程还没有准备好帧，或者已经停止
                    if not self.running: # 相机线程已停止
                        break
                    else:
                        print("⚠️ [主线程] 队列中没有可用帧，等待...或相机线程还在处理。")
                        continue # 继续等待新帧

                # 处理帧
                detections = self.process_frame(frame)
                result_frame = self.draw_results(frame, detections)
                
                cv2.imshow("Real-time Detection", result_frame)
                if cv2.waitKey(1) & 0xFF == ord('q'):
                    print("⏹️ 用户手动停止检测")
                    self.running = False # 停止相机线程
                    break
            
                # 如果视频写入器已初始化，则写入帧
                if self.video_writer and self.video_writer.isOpened():
                    self.video_writer.write(result_frame)

                # 动态调整帧率以匹配 target_processing_fps
                processing_time = time.time() - frame_start_time
                sleep_duration = 0
                if processing_time < (1 / self.target_processing_fps):
                    sleep_duration = (1 / self.target_processing_fps) - processing_time
                    time.sleep(sleep_duration)
                
                # print(f"帧处理时间: {processing_time*1000:.2f}ms, 睡眠时间: {sleep_duration*1000:.2f}ms") # 打印每帧处理时间

                # 更新并打印实时 FPS
                self.frame_count += 1
                current_loop_time = time.time() # 获取当前时间，避免重复调用 time.time()
                if current_loop_time - self.start_time >= 1:
                    self.fps = self.frame_count / (current_loop_time - self.start_time)
                    # print(f"实时 FPS: {self.fps:.2f} (目标: {self.target_processing_fps} FPS)") # 打印实时 FPS 和目标 FPS
                    self.frame_count = 0
                    self.start_time = current_loop_time
                    
        except KeyboardInterrupt:
            print("\n🛑 通过键盘中断停止检测")
            self.running = False # 停止相机线程
        finally:
            self.running = False # 确保停止所有线程
            if self.camera_thread and self.camera_thread.is_alive():
                self.camera_thread.join(timeout=5) # 等待相机线程结束，最多等待5秒
                if self.camera_thread.is_alive():
                    print("⚠️ 相机捕获线程未能正常停止。")

            if self.cap and self.cap.isOpened():
                self.cap.release()
            if self.video_writer and self.video_writer.isOpened(): # 释放视频写入器
                self.video_writer.release()
            cv2.destroyAllWindows()
            print("✅ 资源已释放")

    def _events_to_frame(self, events, width, height):
        """将事件数据转换为彩色图像帧（使用 Metavision SDK 函数）"""
        if events.size == 0:
            return np.zeros((height, width, 3), dtype=np.uint8) # 返回黑色图像
        
        # 使用 viz_events 函数将事件转换为彩色图像
        # viz_events 返回的图像通常是 RGB 格式，OpenCV 默认使用 BGR
        # 所以可能需要转换一下通道顺序
        img_rgb = viz_events(events, height, width)
        img_bgr = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2BGR)
        return img_bgr

    def process_frame(self, frame):
        """处理帧，执行检测"""
        rh, rw = frame.shape[0: 2]
        img_resized = cv2.resize(frame, (self.ckp_resolution, self.ckp_resolution))
        img_rgb = cv2.cvtColor(img_resized, cv2.COLOR_BGR2RGB)
        img_normalized = img_rgb.astype(np.float32) / 255.0
        img_chw = np.transpose(img_normalized, (2, 0, 1))
        img = torch.from_numpy(img_chw).unsqueeze(0)
        
        img = img.to(self.device) # 将图像移动到与模型相同的设备

        model = self.detector.ema.module if self.detector.ema else self.detector.model
        
        if model is None:
            print("❌ 检测模型未正确加载。")
            return np.array([]), np.array([]), np.array([]) # 返回空数组以避免错误

        with torch.no_grad(): # 推理时禁用梯度计算
            outputs = model(img)

        pred_logits = outputs['pred_logits']
        pred_boxes = outputs['pred_boxes']

        prob = F.softmax(pred_logits, dim=-1)
        scores, labels = prob.max(-1)

        results = []
        for i in range(pred_logits.shape[0]):
            score = scores[i]
            label = labels[i]
            box = pred_boxes[i]

            keep = score > self.confidence

            results.append({
                'scores': score[keep],
                'labels': label[keep],
                'boxes': box[keep]
            })
        
        # 提取第一个图像的结果（批量大小为 1）
        scores_final = results[0]['scores'].cpu().numpy()
        labels_final = results[0]['labels'].cpu().numpy()
        boxes_final = results[0]['boxes'].cpu().numpy()

        # 将边界框缩放到原始图像大小
        boxes_final[:, [0, 2]] *= rw
        boxes_final[:, [1, 3]] *= rh
        
        return boxes_final, scores_final, labels_final

    def draw_results(self, frame, detections):
        """在帧上绘制检测结果"""
        boxes, scores, labels = detections
        
        drawn_frame = frame.copy() # 避免直接修改原始帧

        for box, score, cls in zip(boxes, scores, labels):
            cx, cy, w, h = box
            x1 = int(cx - 0.5 * w)
            y1 = int(cy - 0.5 * h)
            x2 = int(cx + 0.5 * w)
            y2 = int(cy + 0.5 * h)

            # 绘制矩形框
            cv2.rectangle(drawn_frame, (x1, y1), (x2, y2), color=(0, 255, 0), thickness=2) # 绿色

            # 获取类别名称
            class_name = mscoco_category2name.get(int(cls), f"未知类别 {int(cls)}")
            label = f"{class_name}: {score:.2f}"

            # 计算文本大小以创建标签的背景矩形
            (text_width, text_height), baseline = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
            cv2.rectangle(drawn_frame, (x1, y1 - text_height - 10), (x1 + text_width, y1), (0, 255, 0), -1) # 绿色背景
            cv2.putText(drawn_frame, label, (x1, y1 - 5),
                        fontFace=cv2.FONT_HERSHEY_SIMPLEX,
                        fontScale=0.5, color=(0, 0, 0), thickness=1, lineType=cv2.LINE_AA) # 黑色文本

        return drawn_frame

if __name__ == "__main__":
    # 配置参数
    current_time_str = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    output_filename = f"detection_output_{current_time_str}.mp4"
    output_directory = "D:\\Propheese-detction\\MvHeatDET\\output"
    save_video_full_path = os.path.join(output_directory, output_filename)

    config = {
        "config_path": "D:\\Propheese-detction\\MvHeatDET\\configs\\evheat\\MvHeatDET.yml",
        # "weights_path": "D:\\Propheese-detction\\MvHeatDET\\best.pth", # 确保这里是 D:\Propheese-detction\MvHeatDET\best.pth
        "weights_path": "D:\\Propheese-detction\\MvHeatDET\\checkpoint0064.pth",
        "camera_source": 0,
        # "camera_source": "D:\\\\Propheese-detction\\\\MvHeatDET\\\\raw\\\\recording_2025-06-14_11-26-35.raw",  # 可以是：
                            # 整数 (如 0, 1) = 传统摄像头索引 (如果 Metavision SDK 不可用)
                            # "" = 默认 Prophesee 事件相机
                            # "/path/to/your/file.raw" 或 "file://path/to/your/file.raw" = Prophesee 事件文件
                            # "rtsp://..." = 网络摄像头
        "frame_width": 1280,
        "frame_height": 720,
        "ckp_resolution": 640, # 将检查点/模型输入分辨率从 640 降低到 480
        "confidence": 0.8,      # 检测置信度阈值
        "save_video_path": save_video_full_path, # 新增保存视频路径
        "target_processing_fps": 8, # 新增目标处理帧率
        "use_amp": True # 启用自动混合精度以加速推理
    }
    
    try:
        detector = RealTimeDetector(config)
        detector.setup_camera(
            src=config["camera_source"],
            width=config["frame_width"],
            height=config["frame_height"]
        )
        detector.run()
    except Exception as e:
        print(f"❌ 程序异常: {str(e)}")